{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Cl2LilHSwHsi"
   },
   "source": [
    "\n",
    "# Objective\n",
    "This notebook demonstrates how blending the capabilities of Elasticsearch as a vector database (VectorDB), search templates, and LLM functions can provide an intelligent query layer.\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/elastic/elasticsearch-labs/blob/main/supporting-blog-content/unifying-elastic-vector-database-and-llms-for-intelligent-query/Unifying_Elastic_Vector_Database_and_LLMs_for_Intelligent_Query.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "\n",
    "- **Elasticsearch as the VectorDB**: Acts as the core search engine, storing and retrieving dense vector embeddings efficiently.\n",
    "- **Search Templates**: Marry index capabilities to query parameters, enabling dynamic query generation and structured search execution.\n",
    "- **LLM Functions**: Parse the possible available parameters within a query and inject them into the search template for a more intelligent and context-aware retrieval process.\n",
    "\n",
    "This combination enables a more sophisticated search experience, leveraging both structured and unstructured data retrieval methods.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kply8eYngIAL"
   },
   "outputs": [],
   "source": [
    "!pip install \"elasticsearch==8.17.0\" openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wcPv-D6lwHsk"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from openai import AzureOpenAI\n",
    "from datetime import datetime, timedelta\n",
    "from zoneinfo import ZoneInfo\n",
    "from elasticsearch import Elasticsearch, helpers, NotFoundError\n",
    "import requests\n",
    "from IPython.display import Markdown, display\n",
    "import requests\n",
    "from getpass import getpass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tWSsNKGOzA8m"
   },
   "source": [
    "##Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9VE-PjkvzESZ"
   },
   "source": [
    "##Elasticsearch Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "U7bQisDzgThO"
   },
   "outputs": [],
   "source": [
    "es_cloud_id = getpass(prompt=\"Enter your Elasticsearch Cloud ID: \")\n",
    "es_api_key = getpass(prompt=\"Enter your Elasticsearch API key: \")\n",
    "\n",
    "es = Elasticsearch(cloud_id=es_cloud_id, api_key=es_api_key, timeout=300)\n",
    "\n",
    "es.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CTQuTPVuwHsk"
   },
   "source": [
    "## Get Completion Method"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MReWDS6VmHih"
   },
   "source": [
    "###Completions Endpoint & API Key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_i3PT0v2lVHI"
   },
   "source": [
    "![2025-02-21_15-48-14.jpg]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D4giT3X4lgbm"
   },
   "outputs": [],
   "source": [
    "ENDPOINT = getpass(\"Azure OpenAI Completions Endpoint: \")\n",
    "\n",
    "AZURE_API_KEY = getpass(\"Azure OpenAI API Key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s6WU8bZ9mFfB"
   },
   "source": [
    "###Deployment Name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OjfotdiRlZDU"
   },
   "source": [
    "![2025-02-21_15-55-02.jpg]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nOLvH3M_ll64"
   },
   "outputs": [],
   "source": [
    "DEPLOYMENT_NAME = getpass(\"Azure OpenAI Deployment Name: \")\n",
    "deployment_name = DEPLOYMENT_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_xVZyR5rmDvS"
   },
   "source": [
    "###API Version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XkJzWUM3lvwO"
   },
   "source": [
    "![2025-02-21_16-01-29.jpg]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vVzZyB8ml0wa"
   },
   "outputs": [],
   "source": [
    "API_VERSION = getpass(\"Completions Endpoint API Version: \")\n",
    "\n",
    "client = AzureOpenAI(\n",
    "    azure_endpoint=ENDPOINT, api_key=AZURE_API_KEY, api_version=API_VERSION\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KRwV4R_dd2FA"
   },
   "source": [
    "##Google Maps API Key Required"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bckurffTnkI4"
   },
   "source": [
    "![2025-02-21_16-25-38.jpg]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dgpGnZQbb7U_"
   },
   "outputs": [],
   "source": [
    "##create google maps api key here: https://developers.google.com/maps/documentation/embed/get-api-key\n",
    "GMAPS_API_KEY = getpass(prompt=\"Enter Google Maps API Key: \")\n",
    "google_maps_api_key = GMAPS_API_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fiVFDWOYUgvH"
   },
   "outputs": [],
   "source": [
    "# Elastic index\n",
    "ES_INDEX = \"hotels\"\n",
    "TEMPLATE_ID = \"hotel_search_template\"\n",
    "\n",
    "# JSON dataset URL\n",
    "DATASET_URL = \"https://ela.st/hotels-dataset\"\n",
    "\n",
    "ELSER_ENDPOINT_NAME = \"my-elser-endpoint\"\n",
    "E5_ENDPOINT_NAME = \"my-e5-endpoint\"\n",
    "\n",
    "\n",
    "# Define the index mapping\n",
    "INDEX_MAPPING = {\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"Address\": {\"type\": \"text\"},\n",
    "            \"Attractions\": {\"type\": \"text\"},\n",
    "            \"Description\": {\"type\": \"text\"},\n",
    "            \"FaxNumber\": {\"type\": \"text\"},\n",
    "            \"HotelCode\": {\"type\": \"long\"},\n",
    "            \"HotelFacilities\": {\"type\": \"text\"},\n",
    "            \"HotelName\": {\"type\": \"text\"},\n",
    "            \"HotelRating\": {\"type\": \"long\"},\n",
    "            \"HotelWebsiteUrl\": {\"type\": \"keyword\"},\n",
    "            \"Map\": {\"type\": \"keyword\"},\n",
    "            \"PhoneNumber\": {\"type\": \"text\"},\n",
    "            \"PinCode\": {\"type\": \"keyword\"},\n",
    "            \"cityCode\": {\"type\": \"long\"},\n",
    "            \"cityName\": {\"type\": \"text\"},\n",
    "            \"combined_fields\": {\n",
    "                \"type\": \"text\",\n",
    "                \"copy_to\": [\"semantic_description_elser\", \"semantic_description_e5\"],\n",
    "            },\n",
    "            \"countryCode\": {\"type\": \"keyword\"},\n",
    "            \"countryName\": {\"type\": \"keyword\"},\n",
    "            \"latitude\": {\"type\": \"double\"},\n",
    "            \"location\": {\"type\": \"geo_point\"},\n",
    "            \"longitude\": {\"type\": \"double\"},\n",
    "            \"semantic_description_e5\": {\n",
    "                \"type\": \"semantic_text\",\n",
    "                \"inference_id\": E5_ENDPOINT_NAME,\n",
    "            },\n",
    "            \"semantic_description_elser\": {\n",
    "                \"type\": \"semantic_text\",\n",
    "                \"inference_id\": ELSER_ENDPOINT_NAME,\n",
    "            },\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "V3mK_FcFy2TA"
   },
   "source": [
    "##Inferencing Endpoint Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RMZoAAq5y0YA"
   },
   "outputs": [],
   "source": [
    "def create_inferencing_endpoints():\n",
    "    endpoints = [\n",
    "        {\n",
    "            \"inference_id\": ELSER_ENDPOINT_NAME,\n",
    "            \"task_type\": \"sparse_embedding\",\n",
    "            \"body\": {\n",
    "                \"service\": \"elasticsearch\",\n",
    "                \"service_settings\": {\n",
    "                    \"num_allocations\": 2,\n",
    "                    \"num_threads\": 1,\n",
    "                    \"model_id\": \".elser_model_2_linux-x86_64\",\n",
    "                },\n",
    "                \"chunking_settings\": {\n",
    "                    \"strategy\": \"sentence\",\n",
    "                    \"max_chunk_size\": 250,\n",
    "                    \"sentence_overlap\": 1,\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "        {\n",
    "            \"inference_id\": E5_ENDPOINT_NAME,\n",
    "            \"task_type\": \"text_embedding\",\n",
    "            \"body\": {\n",
    "                \"service\": \"elasticsearch\",\n",
    "                \"service_settings\": {\n",
    "                    \"num_allocations\": 2,\n",
    "                    \"num_threads\": 1,\n",
    "                    \"model_id\": \".multilingual-e5-small\",\n",
    "                },\n",
    "                \"chunking_settings\": {\n",
    "                    \"strategy\": \"sentence\",\n",
    "                    \"max_chunk_size\": 250,\n",
    "                    \"sentence_overlap\": 1,\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    for endpoint in endpoints:\n",
    "        try:\n",
    "            es.inference.delete(inference_id=endpoint[\"inference_id\"])\n",
    "            print(f\"Deleted endpoint '{endpoint['inference_id']}'\")\n",
    "        except NotFoundError:\n",
    "            print(\n",
    "                f\"Endpoint '{endpoint['inference_id']}' does not exist. Skipping deletion.\"\n",
    "            )\n",
    "\n",
    "        response = es.inference.put(\n",
    "            inference_id=endpoint[\"inference_id\"],\n",
    "            task_type=endpoint[\"task_type\"],\n",
    "            body=endpoint[\"body\"],\n",
    "        )\n",
    "        print(f\"Created endpoint '{endpoint['inference_id']}': {response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QBnpy1HuzK48"
   },
   "source": [
    "##Indexing and Ingestion Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "H3xh9gbezHfO"
   },
   "outputs": [],
   "source": [
    "# Step 1: Create the index with mapping\n",
    "def create_index():\n",
    "    try:\n",
    "        if es.indices.exists(index=ES_INDEX):\n",
    "            print(f\"Index '{ES_INDEX}' already exists. Deleting and recreating...\")\n",
    "            es.indices.delete(index=ES_INDEX)\n",
    "\n",
    "        es.indices.create(index=ES_INDEX, body=INDEX_MAPPING)\n",
    "        print(f\"Index '{ES_INDEX}' created successfully.\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error creating index: {e}\")\n",
    "        exit(1)\n",
    "\n",
    "\n",
    "# Step 2: Download the JSON file\n",
    "def download_json():\n",
    "    print(\"Downloading dataset...\")\n",
    "    print(f\"Using URL: {DATASET_URL}\")\n",
    "\n",
    "    # Start the request\n",
    "    response = requests.get(DATASET_URL, stream=True)\n",
    "    print(f\"Received response with status code: {response.status_code}\")\n",
    "\n",
    "    # Check for errors\n",
    "    try:\n",
    "        response.raise_for_status()\n",
    "        print(\"Response status is OK.\")\n",
    "    except requests.HTTPError as e:\n",
    "        print(f\"HTTP error occurred: {e}\")\n",
    "        raise\n",
    "\n",
    "    # Optionally, show some headers (use carefully in production)\n",
    "    print(\"Response headers:\")\n",
    "    for key, value in response.headers.items():\n",
    "        print(f\"  {key}: {value}\")\n",
    "\n",
    "    # Now return an iterator for the response lines\n",
    "    print(\"Returning line iterator for the response content.\")\n",
    "    return response.iter_lines()\n",
    "\n",
    "\n",
    "# Step 3: Ingest JSON records into Elasticsearch\n",
    "def ingest_data():\n",
    "    print(\"Ingesting data into Elasticsearch...\")\n",
    "    actions = []\n",
    "\n",
    "    for line in download_json():\n",
    "        if line:\n",
    "            record = json.loads(line)\n",
    "            # Convert latitude/longitude to geo_point format\n",
    "            if \"latitude\" in record and \"longitude\" in record:\n",
    "                record[\"location\"] = {\n",
    "                    \"lat\": record[\"latitude\"],\n",
    "                    \"lon\": record[\"longitude\"],\n",
    "                }\n",
    "\n",
    "            actions.append({\"_index\": ES_INDEX, \"_source\": record})\n",
    "\n",
    "            # Bulk index in batches of 50\n",
    "            if len(actions) >= 50:\n",
    "                helpers.bulk(es, actions)\n",
    "                print(f\"Ingested {len(actions)} records...\")\n",
    "                actions = []\n",
    "\n",
    "    # Ingest any remaining records\n",
    "    if actions:\n",
    "        helpers.bulk(es, actions)\n",
    "        print(f\"Ingested {len(actions)} remaining records.\")\n",
    "\n",
    "    print(\"Data ingestion complete.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U9XaRgU8zRDy"
   },
   "source": [
    "##Search Template\n",
    "Removes the existing hotel_search_template if present and replaces it with an updated version. This ensures the template is always current and correctly structured for search operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yiiRn9qdzQBB"
   },
   "outputs": [],
   "source": [
    "# Search template content\n",
    "search_template_content = {\n",
    "    \"script\": {\n",
    "        \"lang\": \"mustache\",\n",
    "        \"source\": \"\"\"{\n",
    "            \"_source\": false,\n",
    "            \"fields\": [\"HotelName\", \"HotelRating\", \"countryName\", \"cityName\", \"countryCode\", \"Attractions\"],\n",
    "            \"retriever\": {\n",
    "                \"standard\": {\n",
    "                    \"query\": {\n",
    "                        \"semantic\": {\n",
    "                            \"field\": \"semantic_description_elser\",\n",
    "                            \"query\": \"{{query}}\"\n",
    "                        }\n",
    "                    },\n",
    "                    \"filter\": {\n",
    "                        \"bool\": {\n",
    "                            \"must\": [\n",
    "                                {{#distance}}{\n",
    "                                    \"geo_distance\": {\n",
    "                                        \"distance\": \"{{distance}}\",\n",
    "                                        \"location\": {\n",
    "                                            \"lat\": {{latitude}},\n",
    "                                            \"lon\": {{longitude}}\n",
    "                                        }\n",
    "                                    }\n",
    "                                }{{/distance}}\n",
    "                                {{#rating}}{{#distance}},{{/distance}}{\n",
    "                                    \"range\": {\n",
    "                                        \"HotelRating\": {\n",
    "                                            \"gte\": {{rating}}\n",
    "                                        }\n",
    "                                    }\n",
    "                                }{{/rating}}\n",
    "                                {{#countryName}}{{#distance}}{{^rating}},{{/rating}}{{/distance}}{{#rating}},{{/rating}}{\n",
    "                                    \"term\": {\n",
    "                                        \"countryName\": \"{{countryName}}\"\n",
    "                                    }\n",
    "                                }{{/countryName}}\n",
    "                                {{#city}}{{#distance}}{{^rating}},{{/rating}}{{/distance}}{{#rating}},{{/rating}}{\n",
    "                                    \"match\": {\n",
    "                                        \"cityName\": \"{{city}}\"\n",
    "                                    }\n",
    "                                }{{/city}}\n",
    "                                {{#countryCode}}{{#distance}}{{^rating}},{{/rating}}{{/distance}}{{#rating}},{{/rating}}{\n",
    "                                    \"term\": {\n",
    "                                        \"countryCode\": \"{{countryCode}}\"\n",
    "                                    }\n",
    "                                }{{/countryCode}}\n",
    "                                {{#distance}}{{^rating}}{{/rating}}{{/distance}}{{#rating}}{{/rating}}\n",
    "                            ],\n",
    "                            \"should\": [\n",
    "                                {{#attraction}}{\n",
    "                                    \"wildcard\": {\n",
    "                                        \"Attractions\": {\n",
    "                                            \"value\": \"*{{attraction}}*\",\n",
    "                                            \"case_insensitive\": true\n",
    "                                        }\n",
    "                                    }\n",
    "                                }{{/attraction}}\n",
    "                            ]\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        }\"\"\",\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "def delete_search_template(template_id):\n",
    "    \"\"\"Deletes the search template if it exists\"\"\"\n",
    "    try:\n",
    "        es.delete_script(id=template_id)\n",
    "        print(f\"Deleted existing search template: {template_id}\")\n",
    "    except Exception as e:\n",
    "        if \"not_found\" in str(e):\n",
    "            print(f\"Search template '{template_id}' not found, skipping delete.\")\n",
    "        else:\n",
    "            print(f\"Error deleting template '{template_id}': {e}\")\n",
    "\n",
    "\n",
    "def create_search_template(\n",
    "    template_id=TEMPLATE_ID, template_content=search_template_content\n",
    "):\n",
    "    \"\"\"Creates a new search template\"\"\"\n",
    "    try:\n",
    "        es.put_script(id=template_id, body=template_content)\n",
    "        print(f\"Created search template: {template_id}\")\n",
    "    except Exception as e:\n",
    "        print(f\"Error creating template '{template_id}': {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qvigzp5w4fTE"
   },
   "source": [
    "## Find A Hotel Method\n",
    "Manages interactions between user queries, LLM functions, and Elasticsearch, orchestrating tool calls to extract search parameters and execute queries.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "L6E7DFWcTEnb"
   },
   "outputs": [],
   "source": [
    "def find_a_hotel(content):\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": (\n",
    "                \"You are an assistant that only provides recommendations \"\n",
    "                \"based on the search results retrieved from Elasticsearch. \"\n",
    "                \"Do not make up information or answer based on assumptions. \"\n",
    "                \"Only use the provided data to respond to the user's queries.\"\n",
    "                \"Don't make assumptions about what values to use with functions. Ask for clarification if a user request is ambiguous.\"\n",
    "            ),\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": content},\n",
    "    ]\n",
    "\n",
    "    tools = [\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"extract_hotel_search_parameters\",\n",
    "                \"description\": \"Extract search parameters for finding hotels (excluding the query itself).  the parameters are extracted from the input query\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"query\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"the full input query\",\n",
    "                        },\n",
    "                        \"distance\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The search radius (e.g., 500m, 1000m).\",\n",
    "                        },\n",
    "                        \"rating\": {\n",
    "                            \"type\": \"number\",\n",
    "                            \"description\": \"The minimum hotel rating (e.g., 3, 4, or 5 stars).\",\n",
    "                        },\n",
    "                        \"location\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Location mentioned in the query (e.g., Belongil Beach, Byron Bay).\",\n",
    "                        },\n",
    "                        \"countryName\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Name of the country (e.g., Australia, Germany).\",\n",
    "                        },\n",
    "                        \"city\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"City name (e.g., Byron Bay, Chicago, Houston).\",\n",
    "                        },\n",
    "                        \"State\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"State or province (e.g., Texas, Alaska, Alberta).\",\n",
    "                        },\n",
    "                        \"countryCode\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The country code (e.g., AU for Australia).\",\n",
    "                        },\n",
    "                        \"attraction\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Hotel attractions, amenities, or descriptive terms (e.g., Beach, Museum, gym, modern, luxurious). This can include multiple options.\",\n",
    "                        },\n",
    "                    },\n",
    "                    \"required\": [\"query\", \"attraction\"],\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"geocode_location\",\n",
    "                \"description\": \"Resolve a location to its latitude and longitude.\",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"location\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The name of the location, e.g., Belongil Beach.\",\n",
    "                        }\n",
    "                    },\n",
    "                    \"required\": [\"location\"],\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "        {\n",
    "            \"type\": \"function\",\n",
    "            \"function\": {\n",
    "                \"name\": \"query_elasticsearch\",\n",
    "                \"description\": \"Query Elasticsearch for accommodations based on provided parameters from extract_hotel_search_parameters.  Must call extract_hotel_search_parameters prior to call this function \",\n",
    "                \"parameters\": {\n",
    "                    \"type\": \"object\",\n",
    "                    \"properties\": {\n",
    "                        \"query\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The original search query (e.g., 'hotels near Belongil Beach').\",\n",
    "                        },\n",
    "                        \"latitude\": {\n",
    "                            \"type\": \"number\",\n",
    "                            \"description\": \"Latitude of the location.\",\n",
    "                        },\n",
    "                        \"longitude\": {\n",
    "                            \"type\": \"number\",\n",
    "                            \"description\": \"Longitude of the location.\",\n",
    "                        },\n",
    "                        \"distance\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"Search radius (e.g., '5000m', '10km').\",\n",
    "                        },\n",
    "                        \"rating\": {\n",
    "                            \"type\": \"number\",\n",
    "                            \"description\": \"Minimum hotel rating (e.g., 3, 4, 5 stars).\",\n",
    "                        },\n",
    "                        \"countryName\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The country name (e.g., 'Australia', 'United States').\",\n",
    "                        },\n",
    "                        \"countryCode\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"The country code (e.g., 'AU', 'US').\",\n",
    "                        },\n",
    "                        \"attraction\": {\n",
    "                            \"type\": \"string\",\n",
    "                            \"description\": \"hotel attractions or amenity (e.g., Beach, Museum, gym, coffee shop, pool). This can be muliple options. Any feature of a hotel can be used here.  Attractions in the query may be obvious so this can be a comprehensive list\",\n",
    "                        },\n",
    "                    },\n",
    "                    \"required\": [\"query\"],\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    parameters = {}\n",
    "    while True:\n",
    "        # Call the LLM with tools\n",
    "        response = client.chat.completions.create(\n",
    "            model=deployment_name,\n",
    "            messages=messages,\n",
    "            tools=tools,\n",
    "            tool_choice=\"auto\",\n",
    "        )\n",
    "\n",
    "        response_message = response.choices[0].message\n",
    "        messages.append(response_message)\n",
    "\n",
    "        # Print formatted messages for debugging\n",
    "        print_messages([response_message])\n",
    "\n",
    "        # Check for tool calls\n",
    "        if response_message.tool_calls:\n",
    "            for tool_call in response_message.tool_calls:\n",
    "                function_name = tool_call.function.name\n",
    "                function_args = json.loads(tool_call.function.arguments)\n",
    "\n",
    "                if function_name == \"extract_hotel_search_parameters\":\n",
    "                    # Debug: Print function_args\n",
    "                    print(\"Function Arguments for extract_hotel_search_parameters:\")\n",
    "                    print(function_args)\n",
    "\n",
    "                    # Extract required and optional parameters\n",
    "                    function_response = handle_extract_hotel_search_parameters(\n",
    "                        function_args\n",
    "                    )\n",
    "\n",
    "                    # Debug: Print function_response\n",
    "                    print(\"Response from handle_extract_hotel_search_parameters:\")\n",
    "                    print(function_response)\n",
    "\n",
    "                    parameters.update(json.loads(function_response))\n",
    "\n",
    "                    # Debug: Print updated parameters\n",
    "                    print(\"Updated parameters after extract_hotel_search_parameters:\")\n",
    "                    print(parameters)\n",
    "\n",
    "                elif function_name == \"query_elasticsearch\":\n",
    "                    # Ensure 'query' is present\n",
    "                    if \"query\" not in parameters:\n",
    "                        print(\"Error: 'query' is required for Elasticsearch queries.\")\n",
    "                        return None\n",
    "\n",
    "                    print(\"Function Arguments for query_elasticsearch:\")\n",
    "                    print(function_args)\n",
    "\n",
    "                    # Update parameters directly\n",
    "                    # parameters.update(function_args)\n",
    "\n",
    "                    # Pass extracted parameters to Elasticsearch\n",
    "                    # function_response = call_elasticsearch(parameters)\n",
    "                    function_response = call_elasticsearch(\n",
    "                        query=function_args.get(\"query\"),\n",
    "                        latitude=function_args.get(\"latitude\"),\n",
    "                        attraction=function_args.get(\"attraction\"),\n",
    "                        longitude=function_args.get(\"longitude\"),\n",
    "                        distance=function_args.get(\"distance\"),\n",
    "                        rating=function_args.get(\"rating\"),\n",
    "                        country_name=function_args.get(\"countryName\"),\n",
    "                        country_code=function_args.get(\"countryCode\"),\n",
    "                    )\n",
    "\n",
    "                elif function_name == \"geocode_location\":\n",
    "                    function_response = geocode_location(\n",
    "                        location=function_args.get(\"location\")\n",
    "                    )\n",
    "                    geo_response = json.loads(function_response)\n",
    "                    parameters.update(geo_response)\n",
    "\n",
    "                    # Debug: Print updated parameters\n",
    "                    print(\"Updated parameters after geocode_location:\")\n",
    "                    print(parameters)\n",
    "                else:\n",
    "                    function_response = json.dumps({\"error\": \"Unknown function\"})\n",
    "\n",
    "                # Append the tool response to the conversation\n",
    "                messages.append(\n",
    "                    {\n",
    "                        \"tool_call_id\": tool_call.id,\n",
    "                        \"role\": \"tool\",\n",
    "                        \"name\": function_name,\n",
    "                        \"content\": json.dumps(function_response),\n",
    "                    }\n",
    "                )\n",
    "        else:\n",
    "            # If no further tools are requested, break the loop\n",
    "            break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6sazS56d4LCT"
   },
   "source": [
    "## Format and Print Messages Functions\n",
    "Formats and prints ChatCompletionMessages for better readability, displaying roles, content, function calls, and tool interactions in a structured way.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S28E_3t5PEib"
   },
   "outputs": [],
   "source": [
    "def format_message(message):\n",
    "    \"\"\"\n",
    "    Format a ChatCompletionMessage for easier readability.\n",
    "    \"\"\"\n",
    "    formatted_message = f\"Role: {message.role}\\n\"\n",
    "    if message.content:\n",
    "        formatted_message += f\"Content: {message.content}\\n\"\n",
    "    if message.function_call:\n",
    "        formatted_message += (\n",
    "            f\"Function Call:\\n\"\n",
    "            f\"  Name: {message.function_call.name}\\n\"\n",
    "            f\"  Arguments: {message.function_call.arguments}\\n\"\n",
    "        )\n",
    "    if message.tool_calls:\n",
    "        formatted_message += \"Tool Calls:\\n\"\n",
    "        for tool_call in message.tool_calls:\n",
    "            formatted_message += (\n",
    "                f\"  Tool Call ID: {tool_call.id}\\n\"\n",
    "                f\"  Function Name: {tool_call.function.name}\\n\"\n",
    "                f\"  Arguments: {tool_call.function.arguments}\\n\"\n",
    "            )\n",
    "    return formatted_message\n",
    "\n",
    "\n",
    "def print_messages(messages):\n",
    "    \"\"\"\n",
    "    Print all ChatCompletionMessages in a nicely formatted way.\n",
    "    \"\"\"\n",
    "    print(\"\\nFormatted Messages:\")\n",
    "    for i, message in enumerate(messages, 1):\n",
    "        print(f\"Message {i}:\")\n",
    "        print(format_message(message))\n",
    "        print(\"-\" * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gyz8L_EL4EHR"
   },
   "source": [
    "## Call Elasticsearch Function\n",
    "Executes a search query using Elasticsearch with structured parameters, leveraging a search template for dynamic query generation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tVWlVx0cNWUd"
   },
   "outputs": [],
   "source": [
    "def call_elasticsearch(\n",
    "    query,\n",
    "    latitude=None,\n",
    "    longitude=None,\n",
    "    distance=None,\n",
    "    rating=None,\n",
    "    country_name=None,\n",
    "    country_code=None,\n",
    "    attraction=None,\n",
    "):\n",
    "    \"\"\"\n",
    "    Call Elasticsearch with the provided parameters and return a JSON-serializable response.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Construct the params dictionary\n",
    "        params = {\n",
    "            \"query\": query,\n",
    "            \"latitude\": latitude,\n",
    "            \"longitude\": longitude,\n",
    "            \"distance\": distance,\n",
    "            \"rating\": rating,\n",
    "            \"countryName\": country_name,\n",
    "            \"countryCode\": country_code,\n",
    "            \"attraction\": attraction,\n",
    "        }\n",
    "\n",
    "        # Remove None values\n",
    "        cleaned_params = {\n",
    "            key: value for key, value in params.items() if value is not None\n",
    "        }\n",
    "\n",
    "        # Debug: Print the parameters for Elasticsearch\n",
    "        print(\"Parameters for Elasticsearch:\")\n",
    "        print(cleaned_params)\n",
    "\n",
    "        # Construct the query body\n",
    "        query_body = {\n",
    "            \"id\": \"hotel_search_template\",  # Replace with your actual template ID\n",
    "            \"params\": cleaned_params,\n",
    "        }\n",
    "\n",
    "        # Debug: Print query for Elasticsearch\n",
    "        print(\"Elasticsearch Query:\")\n",
    "        print(json.dumps(query_body, indent=2))\n",
    "\n",
    "        # Call Elasticsearch\n",
    "        response = es.search_template(index=ES_INDEX, body=query_body)\n",
    "        print(\"Elasticsearch query successful.\")\n",
    "\n",
    "        # Convert response to a JSON-serializable dictionary\n",
    "        response_body = response.body\n",
    "\n",
    "        # Extract and print the number of results\n",
    "        total_results = response_body.get(\"hits\", {}).get(\"total\", {}).get(\"value\", 0)\n",
    "        print(f\"Number of results found: {total_results}\")\n",
    "\n",
    "        return response_body\n",
    "    except Exception as e:\n",
    "        print(f\"Error while querying Elasticsearch: {e}\")\n",
    "        return {\"error\": str(e)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F96sMuQO3_UG"
   },
   "source": [
    "## Handle Extract Hotel Search Parameters Function\n",
    "Validates and processes extracted search parameters, ensuring required values like distance are set for location-based hotel searches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t2E7YSk-LH_J"
   },
   "outputs": [],
   "source": [
    "def handle_extract_hotel_search_parameters(args):\n",
    "    \"\"\"\n",
    "    Validate and handle parameters extracted by the LLM.\n",
    "    \"\"\"\n",
    "    if \"latitude\" in args and \"longitude\" in args:\n",
    "        if \"distance\" not in args:\n",
    "            args[\"distance\"] = \"5000m\"  # Default distance\n",
    "\n",
    "    return json.dumps(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CabxStb63fdK"
   },
   "source": [
    "##Geo Location\n",
    "Resolves a location name into latitude and longitude using Google Geocoding API for geospatial search integration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1EXrB0KdLKH1"
   },
   "outputs": [],
   "source": [
    "def geocode_location(location):\n",
    "    \"\"\"\n",
    "    Resolve a location to latitude and longitude using Google Geocoding API.\n",
    "    \"\"\"\n",
    "    GEOCODING_API_URL = \"https://maps.googleapis.com/maps/api/geocode/json\"\n",
    "    params = {\"address\": location, \"key\": google_maps_api_key}\n",
    "    response = requests.get(GEOCODING_API_URL, params=params)\n",
    "    if response.status_code == 200:\n",
    "        data = response.json()\n",
    "        if data[\"status\"] == \"OK\":\n",
    "            result = data[\"results\"][0][\"geometry\"][\"location\"]\n",
    "            return json.dumps({\"latitude\": result[\"lat\"], \"longitude\": result[\"lng\"]})\n",
    "    return json.dumps({\"error\": \"Geocoding failed\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lYUVqJUS30ij"
   },
   "source": [
    "##Create inferencing endpoints, Index, and Ingestion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jNDaGcTJ2g7g"
   },
   "outputs": [],
   "source": [
    "print(\"Creating inferencing endpoints...\")\n",
    "create_inferencing_endpoints()\n",
    "print(\"Creating hotels index...\")\n",
    "create_index()\n",
    "print(\"Ingesting hotels data...\")\n",
    "ingest_data()\n",
    "print(\"Creating Search Template...\")\n",
    "create_search_template()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Pze4f50gmkn"
   },
   "source": [
    "## Find a hotel\n",
    "Here, issue a query and notice how the query understanding layer matches possible attributes in the query to the attributes supported within the hotels index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OyZKPEPEl4v5",
    "outputId": "f202e5aa-f9a6-4bb5-eccf-90afb6b12bdf"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Formatted Messages:\n",
      "Message 1:\n",
      "Role: assistant\n",
      "Tool Calls:\n",
      "  Tool Call ID: call_VHhjy0TMxPnefibSssrTa5r0\n",
      "  Function Name: extract_hotel_search_parameters\n",
      "  Arguments: {\"query\":\"recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym\",\"distance\":\"250m\",\"rating\":4,\"location\":\"Belongil Beach\",\"attraction\":\"recently renovated, pool, gym\"}\n",
      "\n",
      "--------------------------------------------------\n",
      "Function Arguments for extract_hotel_search_parameters:\n",
      "{'query': 'recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym', 'distance': '250m', 'rating': 4, 'location': 'Belongil Beach', 'attraction': 'recently renovated, pool, gym'}\n",
      "Response from handle_extract_hotel_search_parameters:\n",
      "{\"query\": \"recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym\", \"distance\": \"250m\", \"rating\": 4, \"location\": \"Belongil Beach\", \"attraction\": \"recently renovated, pool, gym\"}\n",
      "Updated parameters after extract_hotel_search_parameters:\n",
      "{'query': 'recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym', 'distance': '250m', 'rating': 4, 'location': 'Belongil Beach', 'attraction': 'recently renovated, pool, gym'}\n",
      "\n",
      "Formatted Messages:\n",
      "Message 1:\n",
      "Role: assistant\n",
      "Tool Calls:\n",
      "  Tool Call ID: call_u1fUp6CBJ8i1p9OZlQMk9kXV\n",
      "  Function Name: geocode_location\n",
      "  Arguments: {\"location\":\"Belongil Beach\"}\n",
      "\n",
      "--------------------------------------------------\n",
      "Updated parameters after geocode_location:\n",
      "{'query': 'recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym', 'distance': '250m', 'rating': 4, 'location': 'Belongil Beach', 'attraction': 'recently renovated, pool, gym', 'latitude': -28.6337328, 'longitude': 153.6003455}\n",
      "\n",
      "Formatted Messages:\n",
      "Message 1:\n",
      "Role: assistant\n",
      "Tool Calls:\n",
      "  Tool Call ID: call_jr4jJ04lW0y0E2AKnVmNumyh\n",
      "  Function Name: query_elasticsearch\n",
      "  Arguments: {\"query\":\"recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym\",\"latitude\":-28.6337328,\"longitude\":153.6003455,\"distance\":\"250m\",\"rating\":4,\"attraction\":\"recently renovated, pool, gym\"}\n",
      "\n",
      "--------------------------------------------------\n",
      "Function Arguments for query_elasticsearch:\n",
      "{'query': 'recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym', 'latitude': -28.6337328, 'longitude': 153.6003455, 'distance': '250m', 'rating': 4, 'attraction': 'recently renovated, pool, gym'}\n",
      "Parameters for Elasticsearch:\n",
      "{'query': 'recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym', 'latitude': -28.6337328, 'longitude': 153.6003455, 'distance': '250m', 'rating': 4, 'attraction': 'recently renovated, pool, gym'}\n",
      "Elasticsearch Query:\n",
      "{\n",
      "  \"id\": \"hotel_search_template\",\n",
      "  \"params\": {\n",
      "    \"query\": \"recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym\",\n",
      "    \"latitude\": -28.6337328,\n",
      "    \"longitude\": 153.6003455,\n",
      "    \"distance\": \"250m\",\n",
      "    \"rating\": 4,\n",
      "    \"attraction\": \"recently renovated, pool, gym\"\n",
      "  }\n",
      "}\n",
      "Elasticsearch query successful.\n",
      "Number of results found: 1\n",
      "\n",
      "Formatted Messages:\n",
      "Message 1:\n",
      "Role: assistant\n",
      "Content: I found one recently renovated accommodation within 250 meters of Belongil Beach that has at least 4 stars and offers both a pool and gym:\n",
      "\n",
      "- **Hotel Name:** Belongil Beach Apartment\n",
      "- **Rating:** 4 stars\n",
      "- **Location:** Byron Bay, New South Wales, Australia\n",
      "\n",
      "If you need more information or assistance with booking, please let me know!\n",
      "\n",
      "--------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "find_a_hotel(\n",
    "    \"recently renovated accommodations 250m from Belongil Beach with at least 4 stars and with a pool and gym\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
